package cn.edu.ncepu;

import java.util.ArrayList;
import java.util.List;

/**
 * 包含扩展欧几里得、丢番图方程、Montgomery模乘算法的核心实现
 */
public class MontgomeryAlgorithm {

    /**
     * 扩展欧几里得算法：求解 ax + by = gcd(a, b)
     * @param a 系数a
     * @param b 系数b
     * @return 数组 [x, y, gcd(a,b)]
     */
    public static long[] extendEuclid(long a, long b) {
        if (a < b) {
            long[] temp = extendEuclid(b, a);
            // 交换x和y，与Python的 return y, x, g 一致
            return new long[]{temp[1], temp[0], temp[2]};
        }

        List<Long> coefficients = new ArrayList<>();
        long originalA = a;
        long originalB = b;

        // 收集a//b的系数
        while (b != 0) {
            coefficients.add(a / b);
            long temp = b;
            b = a % b;
            a = temp;
        }

        // 循环回代计算x, y
        long x = 1, y = 0;
        // 反转系数列表
        java.util.Collections.reverse(coefficients);
        for (long coeff : coefficients) {
            long newX = y;
            long newY = x - y * coeff;
            x = newX;
            y = newY;
        }

        // 返回 [x, y, gcd(a,b)]（最终a即为gcd）
        return new long[]{x, y, a};
    }

    /**
     * 丢番图方程求解：ax + by = c
     * @param a 系数a
     * @param b 系数b
     * @param c 常数项c
     * @return 数组 [x0, y0, xk, yk]，其中：
     *         x = x0 + xk * t
     *         y = y0 + yk * t
     * @throws IllegalStateException 方程无解时抛出
     */
    public static long[] diophantine(long a, long b, long c) {
        long[] eeResult = extendEuclid(a, b);
        long xG = eeResult[0];
        long yG = eeResult[1];
        long g = eeResult[2];

        // 检查是否有解：c必须是g的倍数
        if (c % g != 0) {
            throw new IllegalStateException("No solutions");
        }

        long k = c / g;
        long x0 = xG * k;
        long y0 = yG * k;
        long xk = b / g;
        long yk = -a / g;

        return new long[]{x0, y0, xk, yk};
    }

    /**
     * Montgomery模乘算法实现类
     */
    public static class Montgomery {
        private final long r;    // 基数r（2的幂）
        private final long n;    // 模数n（与r互质）
        private long rPrime;     // r'，满足 r*r' ≡ 1 mod n
        private long nPrime;     // n'，满足 r*r' - n*n' = 1

        /**
         * 构造Montgomery对象，初始化r'和n'
         * @param r 基数r（需与n互质）
         * @param n 模数n
         */
        public Montgomery(long r, long n) {
            this.r = r;
            this.n = n;

            // 求解丢番图方程 r*x + n*y = 1（即 r*x - n*(-y) = 1）
            long[] dioResult = diophantine(r, n, 1);
            this.rPrime = dioResult[0];  // x0
            this.nPrime = dioResult[1];  // y0
            long xk = dioResult[2];      // xk = n/gcd(r,n) = n/1 = n
            long yk = dioResult[3];      // yk = -r/gcd(r,n) = -r

            // 调整r'为正整数
            while (this.rPrime < 0) {
                if (xk > 0) {
                    this.rPrime += xk;
                    this.nPrime += yk;
                    System.out.println(this.rPrime);
                } else {
                    this.rPrime -= xk;
                    this.nPrime -= yk;
                    System.out.println(this.rPrime + " -");
                }
            }
        }

        /**
         * Montgomery约简：将x从Montgomery空间转换回普通空间
         * @param x Montgomery空间的数
         * @return 普通空间的数（x * r' mod n）
         */
        public long reduce(long x) {
            long q = ((x % r) * nPrime) % r;
            long a = (x - q * n) / r;
            // 确保结果为正
            while (a < 0) {
                a += n;
            }
            return a;
        }

        /**
         * 将普通数映射到Montgomery空间
         * @param x 普通空间的数
         * @return Montgomery空间的数（x * r mod n）
         */
        public long present(long x) {
            return (x * r) % n;
        }

        /**
         * 基于Montgomery算法的模幂运算：a^e mod n
         * @param a 底数
         * @param e 指数
         * @param n 模数
         * @return a^e mod n 的结果
         */
        public static long power(long a, long e, long n) {
            // 找到大于n的最小2的幂作为r
            long r = 2;
            while (r <= n) {
                r <<= 1;
            }

            Montgomery montgomery = new Montgomery(r, n);
            long power = 1;
            long result = montgomery.present(1);  // 1映射到Montgomery空间
            long p = montgomery.present(a);       // a映射到Montgomery空间

            // 二进制快速幂核心逻辑
            while (power <= e) {
                if ((power & e) != 0) {
                    // Montgomery空间乘法：result = (result * p) * r' mod n
                    result = montgomery.reduce(result * p);
                }
                // 平方后约简
                p = montgomery.reduce(p * p);
                power <<= 1;  // 等价于power *= 2
            }

            // 将结果从Montgomery空间转换回普通空间
            return montgomery.reduce(result);
        }
    }


}